package com.cx.common.aspect;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.cx.common.annotation.CurrentLimitingAnnotation;
import com.cx.common.constant.RedisConstant;
import com.cx.common.exception.TrafficExceededException;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * @author cx
 * @description: 分布式限流切面
 * @date 2023/7/31 15:02
 */
@Aspect
public class CurrentLimitingAspect {

    public CurrentLimitingAspect(RedisTemplate<String, String> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    private final RedisTemplate<String, String> redisTemplate;

    @Pointcut("@annotation(com.cx.common.annotation.CurrentLimitingAnnotation)")
    public void limitPointcut() {
    }

    @Before("limitPointcut() && @annotation(annotation)")
    public void currentLimiting(CurrentLimitingAnnotation annotation) {
        String key = annotation.key();
        if (StrUtil.isEmpty(key)) {
            throw new RuntimeException("key must is not empty");
        }
        int rate = annotation.rate();
        if (rate <= 0) {
            throw new RuntimeException("rate must is greater than 0");
        }
        int maxSize = annotation.maxSize();
        if (maxSize <= 0) {
            throw new RuntimeException("maxSize must is greater than 0");
        }
        RedisScript<Object> redisScript = new DefaultRedisScript<>(getRedisCurrentLimitingScript(), Object.class);
        List<?> returnResult = (List<?>) redisTemplate.execute(redisScript, Collections.singletonList(key), String.valueOf(rate), String.valueOf(maxSize));
        if (CollUtil.isEmpty(returnResult) || Objects.equals(String.valueOf(returnResult.get(0)), RedisConstant.CUSTOM_FAIL_CODE)) {
            throw new TrafficExceededException("并发已经超过流量限制");
        }
    }

    private String getRedisCurrentLimitingScript() {
        return "local lastQueryTimeKey = KEYS[1]..':'..'time' " +
                "local tokenSizeKey = KEYS[1]..':'..'current' " +
                "local lastQueryTime = tonumber(redis.call('get', lastQueryTimeKey) or '0')  " +
                "local tokenSize = tonumber(redis.call('get', tokenSizeKey) or '0')  " +
                "local rate = tonumber(ARGV[1]) " +
                "local maxSize = tonumber(ARGV[2]) " +
                "local time = redis.call('TIME') " +
                "local now = tonumber(time[1])+tonumber(time[2])/1000000 " +
                "local newTokenSize = (now-lastQueryTime)*rate+tokenSize " +
                "if newTokenSize>=maxSize then newTokenSize = maxSize end " +
                "if newTokenSize > 1 then  " +
                "redis.call('set', lastQueryTimeKey,now) " +
                "return redis.call('set', tokenSizeKey,newTokenSize-1) " +
                "else return 0 " +
                "end ";
    }
}
